import os
from langgraph.types import Command
from typing import TypedDict, Literal
from langgraph.constants import START
from pydantic import BaseModel, Field
from langgraph.graph import StateGraph
from src.common.logger import getLogger
from langchain_core.prompts import ChatPromptTemplate
from src.agentic.agent.multi.StockAgent import StockAgent
from src.agentic.agent.multi.SearchAgent import SearchAgent
from src.agentic.agent.multi.WeatherAgent import WeatherAgent
from src.agentic.agent.multi.CalculateAgent import CalculateAgent

logger = getLogger()

class CollaborationState(TypedDict):
    query: str
    document: str
    answer: str

node_names = ["calculate", "weather", "stock", "search"]

class CollaborationAct(BaseModel):
    agent: Literal[*node_names, "finish"] = Field(description = "根据用户输入问题，判断需要调用的agent，只许返回一个节点")

class CollaborationAgentic:

    def __init__(self, llm_model, agent_tools):
        self.llm_model = llm_model
        self.agent_tool = { tool.name: tool for tool in agent_tools }

    def calculate_node(self, state: CollaborationState) -> Command[Literal[*node_names, "__end__"]]:
        logger.info(f"CollaborationAgentic calculate_node start")
        query = state["query"]
        agent = CalculateAgent(self.llm_model, self.agent_tool["calculate_numerical"])
        response = agent.invoke(query)

        answer = response["answer"]
        next_agent = self.post_process_route(query, answer)
        return Command(goto=next_agent, update={"document": response["document"], "answer": answer})

    def weather_node(self, state: CollaborationState) -> Command[Literal[*node_names, "__end__"]]:
        logger.info(f"CollaborationAgentic weather_node start")
        query = state["query"]
        agent = WeatherAgent(self.llm_model, self.agent_tool["search_weather"])
        response = agent.invoke(query)

        answer = response["answer"]
        next_agent = self.post_process_route(query, answer)
        return Command(goto=next_agent, update={"document": response["document"], "answer": answer})

    def stock_node(self, state: CollaborationState) -> Command[Literal[*node_names, "__end__"]]:
        logger.info(f"CollaborationAgentic stock_node start")
        query = state["query"]
        agent = StockAgent(self.llm_model, self.agent_tool["snatch_stock_price"])
        response = agent.invoke(query)

        answer = response["answer"]
        next_agent = self.post_process_route(query, answer)
        return Command(goto=next_agent, update={"document": response["document"], "answer": answer})

    def search_node(self, state: CollaborationState) -> Command[Literal[*node_names, "__end__"]]:
        logger.info(f"CollaborationAgentic search_node start")
        query = state["query"]
        agent = SearchAgent(self.llm_model, self.agent_tool["search_web"])
        response = agent.invoke(query)

        answer = response["answer"]
        next_agent = self.post_process_route(query, answer)
        return Command(goto = next_agent, update = { "document": response["document"], "answer": answer })

    def post_process_route(self, query, answer):
        route_template = """
            你是一位精准的问题-答案分析专家，请严格按照要求分析并解答，判断当前应归属的节点类型
            
            根据用户问题和答案，返回下列节点中的一个：
            calculate: 用于执行数学计算，只能用于数值计算
            weather: 用于查询城市当天及未来的实时天气情况，禁止用来查询过去时间的天气情况
            stock: 用于查询沪深300指数指定股票的开盘/收盘价
            search: 用于搜索真实世界的信息，如事实、人物、地点等，且不属于上述三类时使用
            finish: 答案不能为空，且满足以下任一条件时必须返回'finish'：
                1、答案包含与用户问题相关的关键词
                2、答案与用户问题语义含义一致
            请直接输出节点名称（如：search），不要附加任何解释或标点，禁止推测或扩展。

            用户问题： {question}
            答案：{answer}
        """
        route_prompt = ChatPromptTemplate.from_template(route_template)
        route_chain = route_prompt | self.llm_model.with_structured_output(CollaborationAct)
        route_result = route_chain.invoke({ "question": query, "answer": answer })
        logger.info(f"CollaborationAgentic post_process_route route_result: {route_result}")
        return route_result.agent

    def build_graph(self):
        logger.info(f"CollaborationAgentic build_graph start")
        graph = StateGraph(CollaborationState)
        graph.add_node("search", self.search_node)
        graph.add_node("calculate", self.calculate_node)
        graph.add_node("weather", self.weather_node)
        graph.add_node("stock", self.stock_node)

        graph.add_edge(START, "search")

        workflow = graph.compile()

        save_path = "D:/Downloads/taixu/images/agentics"
        image_path = os.path.join(save_path, "Agentic_Collaboration_Workflow.png")
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        with open(image_path, 'wb') as file:
            file.write(workflow.get_graph().draw_mermaid_png())

        return workflow

    def invoke(self, query):
        logger.info(f"CollaborationAgentic invoke query: {query}")
        workflow = self.build_graph()
        response = workflow.invoke({ "query": query })
        answer = response.get("answer", None)
        logger.info(f"CollaborationAgentic invoke answer len: {len(answer)}")
        return { "retrieve_docs": response["document"], "chain_result": answer }
